"""
onnx模型推理
"""
import onnxruntime as rt
import cv2
import numpy as np
import cv2
from PIL import Image, ImageDraw, ImageFont
def cv2AddChineseText(img, text, position, textColor=(0, 255, 0), textSize=30):
    if (isinstance(img, np.ndarray)):  # 判断是否OpenCV图片类型
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    # 创建一个可以在给定图像上绘图的对象
    draw = ImageDraw.Draw(img)
    # 字体的格式
    fontStyle = ImageFont.truetype(
        "simsun.ttc", textSize, encoding="utf-8")
    # 绘制文本
    draw.text(position, text, textColor, font=fontStyle)
    # 转换回OpenCV格式
    return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)

# 加载YOLOv5目标检测模型
session = rt.InferenceSession("models/final.onnx")
#input_name = session.get_inputs()[0].name

# 加载类别标签列表
#class_names = ["fall detected", "walking", "sitting", "reflective_clothes", "other_clothes", "hat", "person"]
#class_names = ["other_clothes","reflective_clothes"]
class_names=["reflective_clothes", "other_clothes", "hat", "person"]
# 加载测试图片并记录其大小
img = cv2.imread("test1.jpg")
#height, width, channels = img.shape

# 对测试图片进行预处理，转换为模型所需的输入格式
#img1 = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img1 = cv2.resize(img, (640, 640))
img2 = cv2.resize(img, (640, 640))
img1 = img1.transpose(2, 0, 1)  # HWC -> CHW
img1 = img1.reshape(1, 3, 640, 640).astype('float32') / 255.0
# 执行模型推理，获取检测结果
output = session.run([], {'images': img1})
#output = session.run([], {input_name: img1})
output = np.array(output[0])
detections = output[:, :, :4]
scores = output[:, :, 4:5] * output[:, :, 5:]#
c=output[:, :, 5:]
class_ids = np.argmax(output[:, :, 5:], axis=-1)

# 设置过滤阈值
threshold = 0.8
# 对每个检测结果进行可视化
for i in range(detections.shape[1]):
    x, y, w, h = detections[0,i]
    x1=x-w/2
    y1=y-h/2
    x2=x+w/2
    y2=y+h/2
    class_id = class_ids[0,i]
    score = scores[0,i,class_id]
    if score > threshold :
        print(class_names[class_id])
        print(x,y,w,h)
        print(x1,y1,x2,y2)
        label = class_names[class_id]
        if label=='other_clothes' :
            color = (255,0,255)
            thickness = 1
            font = cv2.FONT_HERSHEY_SIMPLEX
            font_scale = 0.5
            #print('1')
            # 绘制边界框和标签
            cv2.rectangle(img2, (int(x1), int(y1)), (int(x2), int(y2)), color, thickness)
            #print(x1,y1,x2,y2)
            img2=cv2AddChineseText(img2,"未穿反光衣", (int(x1), int(y1) - 5),(0, 255, 0), 20)
        elif label=='person':
            color = (255,0,0)
            thickness = 1
            font = cv2.FONT_HERSHEY_SIMPLEX
            font_scale = 0.5
            #print('1')
            # 绘制边界框和标签
            cv2.rectangle(img2, (int(x1), int(y1)), (int(x2), int(y2)), color, thickness)
            #print(x1,y1,x2,y2)
            #cv2.cv2ImgAddText(img2, f"{'我'}", (int(x1), int(y1) - 5), font, font_scale, color, thickness)
            img2=cv2AddChineseText(img2,"未佩戴安全帽", (int(x1), int(y1) - 5),(0, 255, 0), 20)

# 将可视化后的图像保存到文件或显示在屏幕上
cv2.imwrite("result/result.jpg", img2)
cv2.imshow("Detection Results", img2)
cv2.waitKey(0)


